from flask import Flask, request, jsonify
from sqlalchemy import create_engine, Column, String
from sqlalchemy.orm import sessionmaker, declarative_base
from flask_jwt_extended import JWTManager, jwt_required, create_access_token, get_jwt_identity
from casbin import Enforcer
from casbin_sqlalchemy_adapter import Adapter
from functools import wraps
import uuid

app = Flask(__name__)

# 配置
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///demo.db'
app.config['JWT_SECRET_KEY'] = 'your-secret-key'  # 生产环境请使用安全的密钥
app.config['CASBIN_MODEL_CONF'] = 'rbac_model.conf'

# SQLAlchemy 配置
engine = create_engine(app.config['SQLALCHEMY_DATABASE_URI'], echo=False)
Session = sessionmaker(bind=engine)
Base = declarative_base()

# JWT 初始化
jwt = JWTManager(app)

# Casbin 初始化
adapter = Adapter(engine)
enforcer = Enforcer('rbac_model.conf', adapter)

# 数据库模型
class User(Base):
    __tablename__ = 'users'
    id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    username = Column(String(80), unique=True, nullable=False)
    password = Column(String(120), nullable=False)
    role = Column(String(80), nullable=False)

# 初始化数据库
Base.metadata.create_all(engine)

# 权限检查装饰器
def require_permission():
    def decorator(f):
        @wraps(f)
        @jwt_required()
        def decorated_function(*args, **kwargs):
            user_id = get_jwt_identity()
            resource = kwargs.get('resource')
            action = kwargs.get('action')
            print(resource, action)
            
            session = Session()
            try:
                user = session.query(User).get(user_id)
                if not user:
                    return jsonify({'message': 'User not found'}), 404
                
                print(f"Checking permission: user_id={user_id}, resource={resource}, action={action}")
                if enforcer.enforce(user_id, resource, action):
                    return f(*args, **kwargs)
            finally:
                session.close()
        return decorated_function
    return decorator

# 路由
@app.route('/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    role = data.get('role', 'user')

    session = Session()
    try:
        if session.query(User).filter_by(username=username).first():
            return jsonify({'message': 'User already exists'}), 400

        user = User(username=username, password=password, role=role)
        session.add(user)
        session.commit()

        enforcer.add_role_for_user(user.id, role)
        return jsonify({'message': 'User registered successfully'}), 201
    finally:
        session.close()

@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    session = Session()
    try:
        user = session.query(User).filter_by(username=username, password=password).first()
        if not user:
            return jsonify({'message': 'Invalid credentials'}), 401

        access_token = create_access_token(identity=user.id)
        return jsonify({'access_token': access_token}), 200
    finally:
        session.close()

@app.route('/protected/<resource>/<action>', methods=['GET'])
@require_permission()
def protected(resource, action):
    return jsonify({'message': f'Access granted to {resource} for {action}'}), 200

@app.route('/add_policy', methods=['POST'])
@jwt_required()
def add_policy():
    user_id = get_jwt_identity()
    
    session = Session()
    try:
        user = session.query(User).where(User.id == user_id).first()
        if not user or user.role != 'admin':
            return jsonify({'message': 'Admin access required'}), 403

        data = request.get_json()
        role = data.get('role')
        resource = data.get('resource')
        action = data.get('action')

        enforcer.add_policy(role, resource, action)
        return jsonify({'message': 'Policy added successfully'}), 201
    finally:
        session.close()

if __name__ == '__main__':
    app.run(debug=True)
    # enforcer.add_role_for_user("587f5d5b-2724-48a3-97d6-8793facb70d3", 'user')